package com.gitee.wsl.struct.attribute.pairs

import com.gitee.wsl.math.Math
import kotlin.math.min


/** Analogue of Python's `itertools` with methods needed for AllPairs algorithm.  */
internal object Itertools {
    /**
     * Returns the [Iterator] over `r` length subsequences of elements from the input `List`.
     *
     *
     * Elements are treated as unique based on their position, not on their value.
     * So if the input elements are unique, there will be no repeat values in each combination.
     * <pre>`combinations(Arrays.asList("A", "B", "C", "D"), 2) // --> AB AC AD BC BD CD
     * combinations(Arrays.asList(0, 1, 2, 3), 3) // --> 012 013 023 123
    `</pre> *
     *
     * @param elements the [List] of elements to create combinations from
     * @param r the length of each combination to return
     * @param <T> the type of elements in the [List]
     * @return the [Iterator] over combinations - the `List` of elements of type `T`
    </T> */
    fun <T> combinations(elements: List<T>, r: Int): Iterable<List<T>> {
        return Iterable<List<T>> {
            object : MutableIterator<List<T>> {
                private var currentPosition = 0
                private var currentCombination = IntArray(r)
                private val totalCombinations = calculateCombinationsNumber(elements.size, r)

                private var hasNext = r <= elements.size && r >= 1

                override fun hasNext(): Boolean {
                    return hasNext
                }

                override fun next(): List<T> {
                    if (!hasNext()) {
                        throw NoSuchElementException("No next combination found")
                    }
                    if (currentPosition == 0) {
                        for (i in currentCombination.indices) {
                            currentCombination[i] = i + 1
                        }
                    } else {
                        currentCombination = generateNextCombination(
                            currentCombination,
                            elements.size,
                            r
                        )
                    }
                    val result: MutableList<T> = ArrayList()
                    for (i in currentCombination) {
                        result.add(elements[i - 1])
                    }
                    currentPosition++
                    if (currentPosition >= totalCombinations) {
                        hasNext = false
                    }
                    return result
                }

                override fun remove() {

                }
            }
        }
    }

    /**
     * Returns the [Iterator] over Cartesian product of elements from the input [List]s.
     *
     * @param elementLists the [List] of element [List]s to create Cartesian product from
     * @param <T> the type of elements in the [List]s
     * @return the [Iterator] over products - the [List] of elements of type `T`
    </T> */
    fun <T> product(elementLists: List<List<T>>): Iterable<List<T>> {
        var total: Long = 1
        val max = IntArray(elementLists.size)
        for (i in elementLists.indices) {
            max[i] = elementLists[i].size
        }
        val initProduct = IntArray(elementLists.size)
        initProduct.fill(1)
        for (list in elementLists) {
            total *= list.size.toLong()
        }
        val totalProducts = total
        return Iterable<List<T>> {
            object : MutableIterator<List<T>> {
                private var currentPosition: Long = 0
                private lateinit var currentProduct: IntArray
                override fun hasNext(): Boolean {
                    return currentPosition < totalProducts
                }

                override fun next(): List<T> {
                    if (!hasNext()) {
                        throw NoSuchElementException("No next product found")
                    }
                    if (currentPosition == 0L) {
                        currentProduct = initProduct
                    } else {
                        currentProduct = generateNextProduct(currentProduct, max)
                    }
                    val result: MutableList<T> = ArrayList()
                    for (i in currentProduct.indices) {
                        result.add(elementLists[i][currentProduct[i] - 1])
                    }
                    currentPosition++
                    return result
                }

                override fun remove() {
                    TODO("Not yet implemented")
                }
            }
        }
    }

    private fun calculateCombinationsNumber(n: Int, r: Int): Long {
        var nn = n
        var rr = r
        var nCr: Long = 1
        rr = min(rr.toDouble(), (nn - rr).toDouble()).toInt()
        var i = 1
        while (i <= rr) {
            nCr = if (nn % i == 0) {
                nCr * (nn.toLong() / i)
            } else if (nCr % i == 0L) {
                (nCr / i) * nn
            } else {
                nCr * (nn / i)
            }
            i++
            nn--
        }
        return nCr
    }

    private fun generateNextCombination(currentCombination: IntArray, n: Int, r: Int): IntArray {
        val nextCombination = currentCombination.copyOf(currentCombination.size)
        var nn = n
        var rr = r
        while (nextCombination[rr - 1] == nn) {
            nn--
            rr--
        }
        nextCombination[rr - 1]++
        for (i in rr until r) {
            nextCombination[i] = nextCombination[i - 1] + 1
        }
        return nextCombination
    }

    private fun generateNextProduct(currentProduct: IntArray, max: IntArray): IntArray {
        val nextProduct = currentProduct.copyOf(currentProduct.size)
        val n = nextProduct.size - 1
        nextProduct[n]++
        for (i in n downTo 1) {
            if (nextProduct[i] > max[i]) {
                nextProduct[i] = 1
                nextProduct[i - 1]++
            }
        }
        return nextProduct
    }
}

